from re import T
import os
from transformers import RobertaTokenizer, T5ForConditionalGeneration
from flask import Flask, jsonify, render_template, make_response, json, request
app = Flask(__name__)


test_model = True
if test_model:
    tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base')
    model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base-multi-sum')



def translate_example(code):
    input_ids = tokenizer(code, return_tensors="pt").input_ids
    generated_ids = model.generate(input_ids, max_length=32)
    result = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return result



@app.route("/", methods=['GET','POST'])
def index():
    if request.method == 'GET':
        os.system("rm -rf temp/*")
        return render_template("main.html")
    elif request.method == 'POST':
        return "hahaha"
    
@app.route("/remove", methods=['POST'])
def remove_file():
    os.system("rm -rf temp/*")
    return "done"


@app.route("/upload", methods=['POST'])
def process():
    file = request.files.get('file')
    save_path = os.path.join("temp", file.filename)
    file.save(save_path)
    print("done")

    return "hahaha"


@app.route('/submit', methods=['post'])
def submit():
    data = request.get_json(silent=True)
    print(data['max_len'])
    print(data['lang'])
    filename = os.listdir("temp")
    if len(filename) == 0:
        filename = ""
    else:
        filename = filename[0]
    path = os.path.join("temp", filename)
    
    with open(path, "r", encoding="utf-8") as f:
        data = f.readlines()
    result = []
    id = 1
    with open("result.json", "w") as f:
        for line in data:
            temp_dict = eval(line)
            temp_dict["result"] = translate_example(temp_dict["code"])
            temp_dict["id"] = id
            id += 1
            result.append(temp_dict)
            json.dump(temp_dict, f)
            f.write("\n")
    
    
    return jsonify(result)


@app.route('/submitText', methods=['post'])
def submitText():
    data = request.get_json(silent=True)
    print(data["max_len"])
    print(data['lang'])
    print(data['textarea1'])
    if len(data["textarea1"]) == 0:
        return ""
    
    return translate_example(data['textarea1'])


@app.route('/download', methods=['post'])
def download():
    file = open("result.json", "rb").read()
    response = make_response(file)
    return response




if __name__ == "__main__":
    app.run(debug=True,port= 5002)
    pass